﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.IO;
using System.Diagnostics;

using CVM;

namespace CVML
{
    public class ByteFile
    {
        public class ObjectRecord 
        {
            public byte flags;
            public byte refs = 0xFF;
            public UInt16 index;
            public System.Object adata;

            public void WriteBinary(BinaryWriter bw)
            {
                bw.Write(flags);
                bw.Write(refs);
                bw.Write(index);
                if (adata == null)
                    bw.Write((UInt32)0);
                else if (adata is Single)
                    bw.Write((Single)adata);
                else if (adata is UInt32)
                    bw.Write((UInt32)adata);
                else
                    throw new Exception("ObjectRecord can only hold UInt32 or Single");
            }

            public override string ToString()
            {
                return String.Format("Flags = {0}, Index = {1}, AData = {2}", flags, index, adata);
            }
        }

        KFxn main = new KFxn();
        List<Constant> ks = new List<Constant>();
        Dictionary<Constant, UInt32> odataLookup = new Dictionary<Constant, UInt32>();

        public List<byte> CodeData = new List<byte>();
        public List<ObjectRecord> ObjectData = new List<ObjectRecord>();
        public List<uint> ArrayData = new List<uint>();
        
        public ByteFile(IEnumerable<Term> terms)
        {
            // Gather up all of the terms and constants
            foreach (var t in terms) {
                main.Terms.Add(t);
                if (t is Constant)
                    if (!(t is KInt))
                        AddConstant(t as Constant);
            }
            
            // Create a constant table for the odata
            uint n = 1;
            foreach (var k in ks)
                if (!odataLookup.ContainsKey(k))
                    odataLookup.Add(k, n++ << 1);

            // Create object representations of each constant
            for (int i = 0; i < ks.Count; ++i)
                ObjectData.Add(CreateObject(ks[i]));

            // Iterate over all functions and create a code block
            CreateCodeBlock(CodeData);

            // For each constant with an array component (e.g. KList)
            // I have to set the index in to the main table.
            for (int i = 0; i < ks.Count; ++i)
                SetArrayData(ObjectData[i], ks[i]);
        }

        private void CreateCodeBlock(List<byte> bytes)
        {
            AddCodeBytes(main, bytes);
            foreach (var k in ks)
            {
                if (k is KFxn)
                {
                    var f = k as KFxn;
                    var or = GetObjectRecordForConstant(k);
                    if (bytes.Count >= ushort.MaxValue)
                        throw new Exception("Too much code");
                    or.index = (ushort)bytes.Count;
                    AddCodeBytes(k as KFxn, bytes);
                }
            }
            while (bytes.Count % 4 != 0)
                bytes.Add(0);
        }

        private void AddCodeBytes(KFxn f, List<byte> bytes)
        {
            foreach (var t in f.Terms)
            {
                if (t is Constant)
                {
                    var k = t as Constant;
                    if (k is KInt)
                    {
                        var n = (k as KInt).n;
                        if (n < 0xFF)
                        {
                            bytes.Add((byte)OpCode.Op_byte);
                            bytes.Add((byte)n);
                        }
                        else if (n < 0xFFFF)
                        {
                            bytes.Add((byte)OpCode.Op_short);
                            bytes.Add((byte)(n >> 8));
                            bytes.Add((byte)n);
                        }
                        else 
                        {
                            bytes.Add((byte)OpCode.Op_long);
                            bytes.Add((byte)(n >> 24));
                            bytes.Add((byte)(n >> 16));
                            bytes.Add((byte)(n >> 8));
                            bytes.Add((byte)n);
                        }
                    }
                    else
                    {
                        var n = ConstantToIndex(k);
                        Debug.Assert(n > 0);
                        n >>= 1;
                        n -= 1;
                        if (n < 0xFF)
                        {
                            bytes.Add((byte)OpCode.Op_kbyte);
                            bytes.Add((byte)n);
                        }
                        else if (n < 0xFFFF)
                        {
                            bytes.Add((byte)OpCode.Op_kshort);
                            bytes.Add((byte)(n >> 8));
                            bytes.Add((byte)n);
                        }
                        else
                        {
                            bytes.Add((byte)OpCode.Op_klong);
                            bytes.Add((byte)(n >> 24));
                            bytes.Add((byte)(n >> 16));
                            bytes.Add((byte)(n >> 8));
                            bytes.Add((byte)n);
                        }
                    }
                }
                else if (t is Op)
                {
                    var op = t as Op;
                    bytes.Add((byte)op.Code);
                    for (int i = 0; i < op.Args.Length; ++i)
                        bytes.Add(op.Args[i]);
                }
                else
                {
                    throw new NotImplementedException();
                }
            }
        }

        private void SetArrayData(ObjectRecord o, Constant k)
        {
            if (k is KList)
            {
                KList list = k as KList;
                if (list.Data == null || list.Data.Count == 0)
                    return;

                // Set the index of the adata field of the object
                // One is added because a zero value indicates to no array data.
                o.adata = (uint)(ArrayData.Count + 1);

                // To the adata list we first add the count. 
                ArrayData.Add((uint)list.Data.Count);

                // Then we add an index for each constant
                for (int i = 0; i < list.Data.Count; ++i)
                {
                    var val = list.Data[i];
                    var index = ConstantToIndex(val);
                    ArrayData.Add(index);
                }
            }
            else if (k is KTable)
            {
                // Set the index of the adata field of the object
                // One is added because a zero value indicates to no array data.
                o.adata = (uint)(ArrayData.Count + 1);

                KTable tbl = k as KTable;

                // To the adata list we first add the count. 
                ArrayData.Add((uint)tbl.Keys.Count * 2);

                for (int i = 0; i < tbl.Keys.Count; ++i)
                {
                    var key = tbl.Keys[i] as KInt;
                    if (key == null) throw new Exception("Expected an integer key");
                    uint n = (uint)key.n;
                    n <<= 1;
                    ++n;
                    ArrayData.Add(n);
                    var val = tbl.Values[i];
                    var index = ConstantToIndex(val);
                    ArrayData.Add(index);
                }
            }
        }

        private Constant IndexToConstant(UInt32 n)
        {
            if ((n & 1) != 0)
                return new KInt((int)(n >> 1));
            if (n == 0)
                return null;
            int index = (int)((n >> 1) - 1);
            return ks[index];
        }

        private ObjectRecord GetObjectRecordForConstant(Constant k)
        {
            int n = (int)ConstantToIndex(k);
            Debug.Assert(n % 2 == 0);
            Debug.Assert(n > 1);
            n >>= 1;
            n -= 1;
            return ObjectData[n];
        }

        private UInt32 ConstantToIndex(Constant k)
        {
            if (k is KInt)
            {
                UInt32 r = (UInt32)(k as KInt).n;
                return (r << 1) + 1;
            }
            else
            {
                // Note: the values in k-lookup are one-based indexes
                return odataLookup[k];
            }
        }

        private ObjectRecord CreateObject(Constant k)
        {
            ObjectRecord r = new ObjectRecord();

            if (k is KFloat)
            {
                var flt = k as KFloat;
                r.adata = FloatToUInt32(flt.f);
                r.flags = (byte)TypeFlags.Float << 4;
            }
            else if (k is KFxn)
            {
                var fxn = k as KFxn;                
            }
            else if (k is KList)
            {
                var lst = k as KList;
            }
            else if (k is KTable)
            {
                var tbl = k as KTable;
                r.flags = (byte)TypeFlags.Table << 4;
            }
            else if (k is KInt)
            {
                throw new Exception("Integers are not converted to objects");
            }
            else
            {
                throw new NotImplementedException();
            }

            return r;
        }

        private void AddConstant(Constant k)
        {
            if (k is KInt) return;
            ks.Add(k);
            foreach (var t in k.Children.Where(x => x is Constant))
                AddConstant(t as Constant);
        }

        public IEnumerable<KFxn> Fxns
        {
            get 
            {
                yield return main;
                foreach (var f in from k in ks where k is KFxn select k as KFxn)
                    yield return f;
            }
        }

        public IEnumerable<Term> Terms
        {
            get 
            {
                foreach (var f in Fxns)
                    foreach (var t in f.Terms)
                        yield return t;
            }
        }

        public byte[] GetBytes()
        {
            MemoryStream str = new MemoryStream();
            BinaryWriter bw = new BinaryWriter(str);

            WriteByteArray(CodeData, bw);
            
            // IP
            bw.Write((uint)0);
            
            // Constant data
            bw.Write((uint)ks.Count * 4);
            foreach (var k in ks)
            {
                uint n = ConstantToIndex(k);
                bw.Write(n);
            }

            // Main stack
            bw.Write((uint)0);

            // Auxiliary stack
            bw.Write((uint)0);

            // Object data
            WriteObjectArray(ObjectData, bw);
            
            // Write the array data
            WriteUInt32Array(ArrayData, bw);
            
            return GetBufferCopy(str);
        }

        private void WriteUInt32Array(List<uint> a, BinaryWriter bw)
        {
            bw.Write((uint)(a.Count * 4));
            foreach (var x in a) bw.Write(x);
        }

        private void WriteByteArray(List<byte> a, BinaryWriter bw)
        {
            bw.Write((uint)(a.Count));
            foreach (var x in a) bw.Write(x);
        }

        private void WriteObjectArray(List<ObjectRecord> a, BinaryWriter bw)
        {
            bw.Write((uint)(a.Count * 8));
            foreach (var x in a) x.WriteBinary(bw);
        }

        private byte[] GetBufferCopy(MemoryStream str)
        {
            var r = new byte[str.Position];
            var b = str.GetBuffer();
            for (int i = 0; i < str.Position; ++i)
                r[i] = b[i];
            return r;
        }

        public Constant GetConstantFromPtr(uint p)
        {
            // Check if it is an int.
            if ((p & 1) == 1)
                return new KInt((int)(p >> 1));
            if (p != 0)
                return null;
            p >>= 1;
            p += 1;
            return ks[(int)p];
        }

        public StringBuilder GetCodeRepresentation(StringBuilder sb)
        {
            int i = 0;
            while (i < CodeData.Count)
            {
                int j = i;
                Op op = new Op(CodeData, ref i);
                sb.AppendFormat("{0,3}: {1}\n", j, op);
                if (op.IsTerminator)
                    sb.AppendLine("    ----");
            }
            return sb;
        }

        public StringBuilder GetArrayDataRepresentation(StringBuilder sb)
        {
            sb.AppendLine("Array data");
            sb.AppendFormat("Size = {0}\n", ArrayData.Count);
            int i = 0;
            while (i < ArrayData.Count)
            {
                uint n = ArrayData[i++];
                sb.AppendFormat("Array size = {0}, index = {1}\n", n, i);
                for (int j = 0; j < n; ++j)
                {
                    sb.Append(j).Append(':').Append(ArrayData[i]).Append("; ");
                    ++i;
                }
                sb.AppendLine();
            }
            return sb;
        }

        public StringBuilder GetObjectDataRepresentation(StringBuilder sb)
        {
            sb.AppendLine("Object data");
            sb.AppendFormat("Size = {0}\n", ObjectData.Count);
            for (int i = 0; i < ObjectData.Count; ++i)
                sb.Append(i).Append(':').AppendLine(ObjectData[i].ToString());
            sb.AppendLine();
            return sb;
        }

        private UInt32 FloatToUInt32(float f)
        {
            var ms = new MemoryStream();
            var bw = new BinaryWriter(ms);
            bw.Write(f);
            var br = new BinaryReader(ms);
            return br.ReadUInt32();
        }

        public StringBuilder GetRepresentation(StringBuilder sb)
        {
            GetCodeRepresentation(sb);
            GetObjectDataRepresentation(sb);
            GetArrayDataRepresentation(sb);
            return sb;
        }

        public override string ToString()
        {
            return GetRepresentation(new StringBuilder()).ToString();
        }
    }

    public class ByteFileDisassembler
    {
        public class ObjectRecord
        {
            ByteFileDisassembler data;
            public uint X;
            public uint Y;

            public ObjectRecord(ByteFileDisassembler data, uint x, uint y) 
            {
                this.data = data;
                this.X = x;
                this.Y = y;
            }

            public uint Flags { get { return (X >> 28) & 0xf; } }
            public uint Type { get { return (X >> 24) & 0xf; } }
            public uint Index { get { return X & 0xffff; } }
            public uint Refs { get { return (X >> 16) & 0Xff; } }

            public uint[] Array
            {
                get
                {
                    if (Type == 0 && Y != 0)
                        return data.adata[(int)Y - 1];
                    else
                        return null;
                }
            }

            public override string ToString()
            {
                StringBuilder sb = new StringBuilder();
                sb.Append(Type).Append('%').
                    Append(Flags).Append('%').
                    Append(Index).Append('%').
                    Append(Refs).Append('%').
                    Append(Y);
                return sb.ToString();
            }
        }
        
        public List<byte> code;
        public List<uint> constants;
        public List<uint> stack;
        public List<uint> aux;
        public List<ObjectRecord> odata;
        public List<uint[]> adata;
        public uint ip;

        public ByteFileDisassembler(byte[] bytes)
        {
            MemoryStream ms = new MemoryStream(bytes);
            BinaryReader br = new BinaryReader(ms);

            code = ReadBytes(br);
            ip = br.ReadUInt32();
            constants = ReadList(br);
            stack = ReadList(br);
            aux = ReadList(br);
            odata = ReadObjects(br);
            adata = ReadArrays(br);
        }

        List<byte> ReadBytes(BinaryReader br) {
            var r = new List<byte>();
            uint n = br.ReadUInt32();
            for (int i = 0; i < n; ++i)
                r.Add(br.ReadByte());
            return r;
        }

        List<ObjectRecord> ReadObjects(BinaryReader br) {
            var r = new List<ObjectRecord>();
            var tmp = ReadList(br);
            for (int i = 0; i < tmp.Count; i += 2)
                r.Add(new ObjectRecord(this, tmp[i], tmp[i + 1]));
            return r;
        }

        List<uint[]> ReadArrays(BinaryReader br) {
            var tmp = ReadList(br);
            int i=0; 
            var r = new List<uint[]>();
            while (i < tmp.Count) {
                uint n = tmp[i++];
                var a = new uint[n];
                for (int j=0; j < n; ++j)
                    a[j] = tmp[i++];
                r.Add(a);
            }
            return r;
        }

        List<uint> ReadList(BinaryReader br)
        {
            var r = new List<uint>();
            uint n = br.ReadUInt32();
            Debug.Assert(n % 4 == 0);
            n /= 4;
            for (int i = 0; i < n; ++i)
                r.Add(br.ReadUInt32());
            return r;
        }
    }
}
